package com.pai4j.aigc.llm;
import com.pai4j.aigc.config.properties.LibLibPropertiesConfig;
import com.pai4j.common.constants.ConnectConstants;
import com.pai4j.common.exception.BusinessException;
import com.pai4j.common.helper.SseEmitterHelper;
import com.pai4j.common.util.HttpClientUtil;
import com.pai4j.common.util.IDGenerateUtil;
import com.pai4j.common.util.JsonUtil;
import com.pai4j.domain.vo.llm.image.*;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Service;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
@Service
public class LibLibGenerateImageService extends AbstractGenerateImageService<RequestGenerateParam> {

    @Autowired
    private LibLibPropertiesConfig libLibPropertiesConfig;


    private Map<String, SignatureParam> signatureMap = new HashMap<>();

    /**
     * 线上API域名
     */
    private static final String HOST = "https://openapi.liblibai.cloud";
    /**
     * 旗舰版F.1生图 API
     */
    private static final String GENERATE_API = "/api/generate/webui/text2img/ultra";
    /**
     * 查询结果API
     */
    private static final String RESULT_API = "/api/generate/webui/status";

    private static final String API_PARAM = "?AccessKey=%s&Signature=%s&Timestamp=%s&SignatureNonce=%s";

    @Override
    protected ImageGenerateResponse generate(String userId, RequestGenerateParam param) throws Exception {
        // 提交生图任务
        String taskId = doSubmit(param);
        if (StringUtils.isBlank(taskId)) {
            throw new BusinessException("图片生成失败，请稍后重试～");
        }
        SseEmitterHelper.send(this.getSSEBizCode(), userId, "任务排队中，稍等片刻～");
        ImageGenerateResponse response = null;
        // 查询任务结果
        int count = 1;
        while (response == null) {
            if (count > 15) {
                throw new BusinessException("图片生成失败，请稍后重试～");
            }
            response = getResult(taskId);
            Thread.sleep(2000);
            ++ count;
        }
        SseEmitterHelper.send(this.getSSEBizCode(), userId, "图片已经生成，请查收");
        return response;
    }

    @Override
    protected String doSubmit(RequestGenerateParam param) throws Exception {
        String api = HOST + GENERATE_API + API_PARAM;
        String signatureNonce = IDGenerateUtil.uuid();
        long timestamp = System.currentTimeMillis();
        String signature = this.makeSign(GENERATE_API, libLibPropertiesConfig.getSecretKey(), timestamp, signatureNonce);
        api = String.format(api, libLibPropertiesConfig.getAccessKey(), signature, timestamp, signatureNonce);
        String responseJson = HttpClientUtil.postObjectForUrl(api, param, String.class);
        LibLibGenerateResponse libLibGenerateResponse = JsonUtil.fromJson(responseJson, LibLibGenerateResponse.class);
        if (libLibGenerateResponse == null) {
            return null;
        }
        if (ObjectUtils.equals(libLibGenerateResponse.getCode(), 0)) {
            String taskId = libLibGenerateResponse.getData().getGenerateUuid();
            signatureMap.put(taskId, new SignatureParam(signatureNonce, timestamp, signature));
            return taskId;
        }
        return null;
    }

    @Override
    protected ImageGenerateResponse getResult(String taskId) throws Exception {

        String api = HOST + RESULT_API + API_PARAM;
        SignatureParam signatureParam = signatureMap.get(taskId);
        String signatureNonce = signatureParam.getSignatureNonce();
        long timestamp = signatureParam.getTimestamp();
        String signature = this.makeSign(RESULT_API, libLibPropertiesConfig.getSecretKey(), timestamp, signatureNonce);
        api = String.format(api, libLibPropertiesConfig.getAccessKey(), signature, timestamp, signatureNonce);
        Map<String, Object> body = new HashMap<>();
        body.put("generateUuid", taskId);
        Map<String, String> header = new HashMap<>();
        header.put("Content-Type", "application/json");
        String responseJson = HttpClientUtil.postObjectForUrl(api, body, header, String.class);

        LibLibGenImageResponse libLibGenImageResponse = JsonUtil.fromJson(responseJson, LibLibGenImageResponse.class);
        if (libLibGenImageResponse != null && ObjectUtils.equals(libLibGenImageResponse.getCode(), 0)) {
            LibLibGenImageResponseData data = libLibGenImageResponse.getData();
            if (data == null) {
                return null;
            }
            List<LibLibGenImageResponseData.LibLibGenImageResponseDataImage> images = data.getImages();
            if (CollectionUtils.isEmpty(images)) {
                return null;
            }
            ImageGenerateResponse response = new ImageGenerateResponse();
            List<String> imageUrls = images.stream().map(LibLibGenImageResponseData.LibLibGenImageResponseDataImage::getImageUrl).collect(Collectors.toList());
            response.setImages(imageUrls);
            response.setCode(HttpStatus.OK.value());
            return response;
        }
        return null;
    }

    public String makeSign(String api, String sk, long timestamp, String signatureNonce) {
        // 拼接请求数据
        String content = api + "&" + timestamp + "&" + signatureNonce;
        try {
            // 生成签名
            SecretKeySpec secret = new SecretKeySpec(sk.getBytes(), "HmacSHA1");
            Mac mac = Mac.getInstance("HmacSHA1");
            mac.init(secret);
            return Base64.encodeBase64URLSafeString(mac.doFinal(content.getBytes()));
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("no such algorithm");
        } catch (InvalidKeyException e) {
            throw new RuntimeException(e);
        }
    }

    private String getSSEBizCode() {
        return ConnectConstants.ConnectBiz.getBizCode(ConnectConstants.ClientEnum.PC,
                ConnectConstants.AppEnum.WEB,
                ConnectConstants.BizEnum.CHAT);
    }


    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    public static class SignatureParam {

        String signatureNonce;

        long timestamp;

        String signature;
    }
}
